#include "main.h"

//#define CUSTOM

extern  App *theApp;

static inline int
get_send_comp(struct rdma_cm_id *id, struct ibv_wc *wc)
{
#ifdef CUSTOM
	struct ibv_cq *cq;
	int ret;

	do {
		ret = ibv_poll_cq(id->send_cq, 1, wc);
		if (ret)
			break;
	} while (1);

	return (ret < 0) ? rdma_seterrno(ret) : ret;
#else
	return rdma_get_send_comp(id, wc);
#endif
}

static inline int
get_recv_comp(struct rdma_cm_id *id, struct ibv_wc *wc)
{
#ifdef CUSTOM
	struct ibv_cq *cq;
	void *context;
	int ret;

	do {
		ret = ibv_poll_cq(id->recv_cq, 1, wc);
		if (ret)
			break;
	} while (1);

	return (ret < 0) ? rdma_seterrno(ret) : ret;
#else
	return rdma_get_recv_comp(id, wc);
#endif
}


RDMAUtil::RDMAUtil(ep_mode mode, int port)
{
	rdmaBuffer = (byte*)malloc(MR_LEN);

	ec = rdma_create_event_channel();
	if (!ec) {
		LOG("Failed to create event channel\n");
	 	exit(1);
	}

	int ret = rdma_create_id(ec, &id, NULL, RDMA_PS_TCP);
	if (ret < 0) {
		LOG("Failed to create id\n");
	 	exit(1);
	}

	if (mode == EP_CLIENT) {
		initClient(port);
	} else if (mode == EP_SERVER) {
		initServer(port);
	}
}

RDMAUtil::~RDMAUtil()
{
	ibv_destroy_comp_channel(id->send_cq_channel);
	ibv_destroy_cq(id->send_cq);
	ibv_dealloc_pd(id->pd);
	rdma_destroy_qp(id);
	rdma_destroy_id(id);
	ibv_dereg_mr(mr);
	free(rdmaBuffer);
}

struct ibv_context* RDMAUtil::get_context()
{
	int i, ret = 0, num_devices;
    struct ibv_device *ib_dev = NULL, **dev_list;
	struct ibv_context* ibv_ctx;
	struct ibv_device_attr dev_attr;

	// get IB device list
    dev_list = ibv_get_device_list(&num_devices);
	if (!dev_list) {
		LOG("Failed to get IB devices list\n");
		ret = -1;
		goto free_device_list;
	}

	if (!num_devices) {
		LOG("No IB device found\n");
		ret = -1;
		goto free_device_list;
	}

	// get specific IB device
	for (i = 0; i < num_devices; i++) {
		string dev_name = string(ibv_get_device_name(dev_list[i]));
		if (gConfig->rdmaDeviceName == "") {
			gConfig->rdmaDeviceName = dev_name;
		}
		if (dev_name == gConfig->rdmaDeviceName) {
			ib_dev = dev_list[i];
			break;
		}
	}

	// device not found	
	if (!ib_dev) {
		LOG("IB device %s not found\n", gConfig->rdmaDeviceName.c_str());
		ret = -1;
		goto free_device_list;
	}

	// open device to get ibv_context
	ibv_ctx = ibv_open_device(ib_dev);
	if (!ibv_ctx) {
		LOG("Failed to open IB device %s\n", gConfig->rdmaDeviceName.c_str());
		ret = -1;
		goto close_device;
	}

	// query device attribute
	ret = ibv_query_device(ibv_ctx, &dev_attr);
	if (ret < 0) {
		LOG("Failed to query device attribute\n");
		goto close_device;
	}

	ibv_free_device_list(dev_list);
	return ibv_ctx;

close_device:
	if (!ibv_ctx)
		ibv_close_device(ibv_ctx);

free_device_list:
	if(!dev_list) 
		ibv_free_device_list(dev_list);

	return NULL;
}

int RDMAUtil::build_ep()
{
	int ret, num_wr;
	struct ibv_pd *pd;
	struct ibv_cq *cq;
	struct ibv_qp *qp;
	struct ibv_comp_channel *cc;
	struct ibv_qp_init_attr init_attr;
	struct rdma_addrinfo hints, *res;

	if (!id->verbs){
		LOG("id->verbs is NULL!\n");
		id->verbs = get_context();
	}

	// alloc protection domain
	if(!id->pd){
		pd = ibv_alloc_pd(id->verbs);
		if (!pd) {
			LOG("Failed to alloc protection domain");
			ret = -1;
			goto dealloc_pd;
		}
		id->pd = pd;
	}

	// create completion channel
	cc = ibv_create_comp_channel(id->verbs);
	if (!cc) {
		LOG("Failed to create completion channel\n");
		ret = -1;
		goto destroy_cc;
	}

	// create completion queue
	cq = ibv_create_cq(id->verbs, MAX_MR, NULL, cc, 0);
	if (!cq) {
		LOG("Failed to create completion queue\n");
		ret = -1;
		goto destroy_cq;
	}

	// set queue pair init attribute
	memset(&init_attr, 0, sizeof(init_attr));
	init_attr.sq_sig_all = 1;
	init_attr.qp_type = IBV_QPT_RC;
	init_attr.send_cq = cq;
	init_attr.recv_cq = cq;
	init_attr.cap.max_send_wr = init_attr.cap.max_recv_wr = MAX_MR;
	init_attr.cap.max_send_sge = init_attr.cap.max_recv_sge = MAX_SGE;
	

	// create queue pair
	ret = rdma_create_qp(id, id->pd, &init_attr);
	if (ret < 0) {
		ret = -1;
		goto destroy_qp;
	}

	cq->cq_context = id;
	id->send_cq = cq;
	id->recv_cq = cq;
	id->send_cq_channel = cc;
	id->recv_cq_channel = cc;
	return 0;

destroy_qp:
	rdma_destroy_qp(id);

destroy_cq:
	if (!cq)
		ibv_destroy_cq(cq);

destroy_cc:
	if (!cc)
		ibv_destroy_comp_channel(cc);

dealloc_pd:
	if (!pd)
		ibv_dealloc_pd(pd);

	return ret;
}

void RDMAUtil::initClient(int port)
{
	int ret = 0;
  	struct addrinfo *res;
	struct rdma_cm_event *event;

	string addr = gConfig->rdmaInetAddress;
	char *name, *service;

	if (addr == "")	addr = gConfig->outputAddresses[theApp->index];
	
	name = const_cast<char*> (addr.c_str());
	service = const_cast<char*> (toString(port).c_str());

	ret = getaddrinfo(name, service, NULL, &res);
	if (ret < 0) {
		LOG("Failed to get addrinfo\n");
	 	exit(1);
	}

	ret = rdma_resolve_addr(id, NULL, res->ai_addr, 500);
	if (ret < 0) {
		LOG("Failed to resolve addr\n");
	 	exit(1);
	}

	while (rdma_get_cm_event(ec, &event) == 0) {
		struct rdma_cm_event event_copy;
		memcpy(&event_copy, event, sizeof(*event));
		rdma_ack_cm_event(event);

		if (event_copy.event == RDMA_CM_EVENT_ADDR_RESOLVED){
			ret = rdma_resolve_route(id, 500);
			if (ret < 0) {
				LOG("Failed to resolve route\n");
				exit(1);
			}
		} else if (event_copy.event == RDMA_CM_EVENT_ROUTE_RESOLVED){
			ret = build_ep();
			if (ret < 0) {
				LOG("Failed to build endpoint\n");
				exit(1);
			}

			mr = ibv_reg_mr(id->pd, rdmaBuffer, MR_LEN, IBV_ACCESS_LOCAL_WRITE);
			if (!mr) {
				LOG("Failed to register memmory\n");
				exit(1);
			}

			ret = rdma_connect(id, NULL);
			if (ret < 0) {
				LOG("Failed to connect %s\n", strerror(errno));
				exit(1);
			}
		} else if(event_copy.event == RDMA_CM_EVENT_ESTABLISHED){
			LOG("Connected to remote pipeline on %s:%d\n", addr.c_str(), port);		
			break;
		}
  	}
}

void RDMAUtil::initServer(int port)
{
	int ret;
  	struct sockaddr_in addr;
	struct rdma_cm_event *event;

	memset(&addr, 0, sizeof(addr));
	addr.sin_family = AF_INET;
	addr.sin_port = htons(port);

	ret = rdma_bind_addr(id, (struct sockaddr *)&addr);
	if (ret < 0) {
		LOG("Failed to bind addr\n");
		exit(1);
	}
	
	ret = rdma_listen(id, 1);
	if (ret < 0) {
		LOG("Failed to listen on server\n");
		exit(1);
	}

	LOG("Waiting for rdma connection on port %d\n", port);

	while (rdma_get_cm_event(ec, &event) == 0) {
		struct rdma_cm_event event_copy;
		memcpy(&event_copy, event, sizeof(*event));
		rdma_ack_cm_event(event);

		if (event_copy.event == RDMA_CM_EVENT_CONNECT_REQUEST){
			id = event_copy.id;
			ret = build_ep();
			if (ret < 0) {
				LOG("Failed to build endpoint\n");
				exit(1);
			}

			mr = ibv_reg_mr(id->pd, rdmaBuffer, MR_LEN, IBV_ACCESS_LOCAL_WRITE);
			if (!mr) {
				LOG("Failed to register memmory\n");
				exit(1);
			}

			ret = rdma_accept(id, NULL);
			if (ret < 0) {
				LOG("Failed to accept\n");
				exit(1);
			}
		} else if(event_copy.event == RDMA_CM_EVENT_ESTABLISHED){
			LOG("Connection established\n");
			break;
		}
  	}

}

int RDMAUtil::send_length(int length)
{
	int ret;
	struct ibv_wc wc;
	memcpy(rdmaBuffer, &length, sizeof(int));
	
	rdma_post_send(id, NULL, rdmaBuffer, sizeof(int), mr, IBV_SEND_INLINE);
	if (ret < 0) {
		LOG("Failed to post send\n");
		return -1;
	}

	ret = get_send_comp(id, &wc);
	if (ret < 0 || wc.status > 0) {
		LOG("Failed to get send comp\n");
		return -1;
	}

	return 0;
}

int RDMAUtil::recv_length()
{
	int ret;
	struct ibv_wc wc;

	ret = rdma_post_recv(id, NULL, rdmaBuffer, sizeof(int), mr);
	if (ret < 0) {
		LOG("Failed to post recv\n");
		return -1;
	}

	ret = get_recv_comp(id, &wc);
	if (ret < 0 || wc.status > 0) {
		LOG("Failed to get recv comp\n");
		return -1;
	}

	memcpy(&ret, rdmaBuffer, sizeof(int));
	return ret;
}

int RDMAUtil::send(byte *buffer, int bufLen)
{
	int ret;
	struct ibv_wc wc;

	int num_wr = (bufLen + MR_LEN -1) / MR_LEN;
	int last_wr_len = bufLen % MR_LEN;
	last_wr_len = (last_wr_len == 0) ? MR_LEN : last_wr_len;

	for (int i = 0; i < num_wr; i++) {
		int len = (i+1 == num_wr) ? last_wr_len : MR_LEN;
		memcpy(rdmaBuffer, buffer+i*MR_LEN, len);

		//TODO: signal once
		ret = rdma_post_send(id, NULL, rdmaBuffer, len, mr, 0);
		if (ret < 0) {
			LOG("Failed to post send\n");
			return -1;
		}

		ret = get_send_comp(id, &wc);
		if (ret < 0 || wc.status > 0) {
			LOG("Failed to get send comp\n");
			return -1;
		}
	}

	return 0;
}


int RDMAUtil::recv(byte *buffer, int bufLen)
{
	int ret;
	struct ibv_wc wc;

	int num_wr = (bufLen + MR_LEN -1) / MR_LEN;
	int last_wr_len = bufLen % MR_LEN;
	last_wr_len = (last_wr_len == 0) ? MR_LEN : last_wr_len;

	for (int i = 0; i < num_wr; i++) {
		int len = (i+1 == num_wr) ? last_wr_len : MR_LEN;

		//TODO: signal once
		ret = rdma_post_recv(id, NULL, rdmaBuffer, len, mr);
		if (ret < 0) {
			LOG("Failed to post recv\n");
			return -1;
		}

		ret = get_recv_comp(id, &wc);
		if (ret < 0 || wc.status > 0) {
			LOG("Failed to get recv comp\n");
			return -1;
		}

		memcpy(buffer+i*MR_LEN, rdmaBuffer, len);
	}

	return bufLen;
}

int RDMAUtil::send2(byte *buffer, int bufLen)
{
	if(bufLen < 4096){
		return send(buffer, bufLen);
	} else {
		int ret;
		struct ibv_wc wc;

		rdmaBuffer[0] = '*';
		memcpy(rdmaBuffer+1, &bufLen, sizeof(int));
	
		ret = rdma_post_send(id, NULL, rdmaBuffer, sizeof(int)+1, mr, IBV_SEND_INLINE);
		if (ret < 0) {
			LOG("Failed to post send\n");
			return -1;
		}

		ret = get_send_comp(id, &wc);
		if (ret < 0 || wc.status > 0) {
			LOG("Failed to get send comp\n");
			return -1;
		}

		return send(buffer, bufLen);
	}
}

int RDMAUtil::recv2(byte *buffer)
{
	int ret;
	struct ibv_wc wc;

	ret = rdma_post_recv(id, NULL, rdmaBuffer, MR_LEN, mr);
	if (ret < 0) {
		LOG("Failed to post recv\n");
		return -1;
	}

	ret = get_recv_comp(id, &wc);
	if (ret < 0 || wc.status > 0) {
		LOG("Failed to get recv comp\n");
		return -1;
	}

	if(wc.byte_len == (sizeof(int) + 1) && rdmaBuffer[0] == '*'){ 
		memcpy(&ret, rdmaBuffer+1, sizeof(int));
		if(ret > 0) return recv(buffer, ret);
		else return ret;
	} else {
		memcpy(buffer, rdmaBuffer, wc.byte_len);
		return wc.byte_len;
	}
}